#include"vibe3.h"

__device__ inline int safe_max(int a, int b) 
{
    return (a > b) ? a : b;
}

__device__ inline int safe_min(int a, int b) 
{
    return (a < b) ? a : b;
}

// 1080*1920*20 

__global__ void random_setup_kernel(curandStatePhilox4_32_10_t *d_state, unsigned long long seed, int H, int W) 
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    int idy = blockDim.y * blockIdx.y + threadIdx.y;

    if (idx < W && idy < H) 
    {
        int state_idx = idx + idy * W;
        curand_init(seed, state_idx, 0, &d_state[state_idx]);
    }
}

void random_setup(curandStatePhilox4_32_10_t *d_state, unsigned long long seed, int H, int W,int BLOSIZE,cudaStream_t stream)
{
    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    random_setup_kernel<<<block_per_grid, thread_per_block, 0, stream>>>(d_state, time(NULL), H, W);
    cudaStreamSynchronize(stream);
    // 错误检查
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
    }
}





__global__ void ProcessFirstFrame_kernel(float* gray, float* samples, curandStatePhilox4_32_10_t* d_state, const int* offset, int num_samples, int H, int W) {
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    int idy = blockDim.y * blockIdx.y + threadIdx.y;

    if (idx < W && idy < H) {
        for (int k = 0; k < num_samples; ++k) {
            int state_idx = idx + idy * W;
            float random_float1 = curand_uniform(&d_state[state_idx]);
            int random_int1 = (int)(random_float1 * 9);
            float random_float2 = curand_uniform(&d_state[state_idx]);
            int random_int2 = (int)(random_float2 * 9);
            // printf("%d %d\n",random_int1,random_int2);
            int pos_idx = (idy + offset[random_int1]) * W + (idx + offset[random_int2]);
            pos_idx = safe_max(0, safe_min(H * W - 1, pos_idx));
            samples[k * H * W + idy * W + idx] = gray[pos_idx];
        }
    }
}

void ProcessFirstFrame(float* gray, float* samples, curandStatePhilox4_32_10_t* d_state, const int* offset, int num_samples, int H, int W, int BLOSIZE,cudaStream_t stream)
{
    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    ProcessFirstFrame_kernel<<<block_per_grid, thread_per_block, 0, stream>>>(gray, samples,d_state, offset,num_samples,H, W);
    cudaStreamSynchronize(stream);
    // 错误检查
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
    }
}



__global__ void ProcessNextFrame_kernel(float* gray, float* samples, int* img_count, int* target_hit_count, int* d_mask_img, curandStatePhilox4_32_10_t* d_state, int* offset, int H, int W, int min_match, int random_sample, int num_samples, int maxhitcount) 
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    int idy = blockDim.y * blockIdx.y + threadIdx.y;

    if (idx < W && idy < H) 
    {
        int thread_pos_idx = idy * W + idx;
        for (int k = 0; k < num_samples; ++k) 
        {
            if (abs(gray[thread_pos_idx] - samples[k * H * W + idy * W + idx]) < 15) 
            {
                img_count[thread_pos_idx] += 1;
            }
            if (img_count[thread_pos_idx] == min_match) 
            {
                break;
            }
        }
        if (img_count[thread_pos_idx] >= min_match) 
        {
            d_mask_img[thread_pos_idx] = 0;
            target_hit_count[thread_pos_idx] = 0;
            
            int random_int3 = (int)(curand_uniform(&d_state[thread_pos_idx]) * random_sample);
            if (random_int3 == 0) 
            {
                int random_int4 = (int)(curand_uniform(&d_state[thread_pos_idx]) * 9);
                int random_int5 = (int)(curand_uniform(&d_state[thread_pos_idx]) * 9);
                int pos_idx = (idy + offset[random_int4]) * W + (idx + offset[random_int5]);
                pos_idx = safe_max(0, safe_min(H * W - 1, pos_idx));
                int random_int6 = (int)(curand_uniform(&d_state[thread_pos_idx]) * num_samples);
                samples[random_int6 * H * W + idy * W + idx] = gray[pos_idx];
            }
            int random_int7 = (int)(curand_uniform(&d_state[thread_pos_idx]) * random_sample);
            if (random_int7 == 0) 
            {
                int random_int4 = (int)(curand_uniform(&d_state[thread_pos_idx]) * 9);
                int random_int5 = (int)(curand_uniform(&d_state[thread_pos_idx]) * 9);
                int pos_idx = (idy + offset[random_int4]) * W + (idx + offset[random_int5]);
                pos_idx = safe_max(0, safe_min(H * W - 1, pos_idx));
                int random_int8 = (int)(curand_uniform(&d_state[thread_pos_idx]) * num_samples);
                samples[random_int8 * H * W + idy * W + idx] = gray[pos_idx];
            }
        } 
        else 
        {
            d_mask_img[thread_pos_idx] = 255;
            target_hit_count[thread_pos_idx] += 1;
            if (target_hit_count[thread_pos_idx] >= maxhitcount) 
            {
                int random_int9 = (int)(curand_uniform(&d_state[thread_pos_idx]) * num_samples);
                samples[random_int9 * H * W + idy * W + idx] = gray[thread_pos_idx];
            }
            
        }
    }
}



void ProcessNextFrame(float* gray, float* samples, int* img_count, int* target_hit_count, int* d_mask_img, curandStatePhilox4_32_10_t* d_state, int* offset, int H, int W, int min_match, int random_sample, int num_samples, int maxhitcount,int BLOSIZE,cudaStream_t stream)
{
    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    ProcessNextFrame_kernel<<<block_per_grid,thread_per_block, 0, stream>>>(gray, samples,img_count, target_hit_count, d_mask_img,d_state, offset, H, W, min_match,random_sample, num_samples, maxhitcount);
    cudaStreamSynchronize(stream);
    // 错误检查
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
    }
}

__global__ void Init_Matarix_kernel(float* samples,int* img_count,int* target_hit_count,int* img_mask,int num_samples,int H,int W)
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    int idy = blockDim.y * blockIdx.y + threadIdx.y;
    if (idx<W && idy<H)
    {
        img_count[idx + idy * W] = 0;
        img_mask[idx + idy * W] = 0;
        target_hit_count[idx + idy * W] = 0;
        for(int m=0;m<num_samples;m++)
        {
            samples[m * H * W + idx + idy * W] = 0;
        }
    }
}

void Init_Matarix(float* samples,int* img_count,int* target_hit_count,int* img_mask,int num_samples,int H,int W,int BLOSIZE,cudaStream_t stream)
{
    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    Init_Matarix_kernel<<<block_per_grid,thread_per_block, 0, stream>>>( samples,img_count,target_hit_count,img_mask,num_samples,H,W);
    cudaStreamSynchronize(stream);
    // 错误检查
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
    }
}


__global__ void Reset_Matarix_kernel(int*img_count,int* img_mask,int H,int W)
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    int idy = blockDim.y * blockIdx.y + threadIdx.y;
    if (idx<W && idy<H)
    {
        img_count[idx + idy * W] = 0;
        // img_mask[idx + idy * W] = 0;
    }
}

void Reset_Matarix(int*img_count,int* img_mask,int H,int W,int BLOSIZE,cudaStream_t stream)
{
    dim3 thread_per_block(BLOSIZE, BLOSIZE);
    dim3 block_per_grid((W + BLOSIZE - 1) / BLOSIZE, (H + BLOSIZE - 1) / BLOSIZE);
    Reset_Matarix_kernel<<<block_per_grid,thread_per_block, 0, stream>>>(img_count,img_mask,H,W);
    cudaStreamSynchronize(stream);
    // 错误检查
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        fprintf(stderr, "CUDA error: %s\n", cudaGetErrorString(err));
    }
}